from typing import List, Dict, Any, Tuple, Optional
from .scheduler_base import SchedulerBase

class PSSScheduler(SchedulerBase):
    """Progressive Synchronous Scheduler implementation.
    
    This scheduler implements the PSS strategy for balanced task distribution
    based on worker performance and task time costs.
    """
    
    def __init__(self, config: Dict[str, Any]):
        super().__init__(config)
        self.num_sync_points = config['scheduler']['sync_points']
        self.load_balance = config['scheduler']['load_balance']
        self.worker_loads = {}
        self.task_segments = []
        self.worker_speeds = {}  # Worker processing speeds
        self.task_costs = {}    # Estimated time cost for each task
        
    def initialize(self) -> bool:
        """Initialize the PSS scheduler.
        
        Sets up task segments and initializes worker performance tracking.
        
        Returns:
            bool: True if initialization succeeds
        """
        try:
            # Initialize worker tracking
            for worker in self.config['system']['workers']:
                worker_id = worker['host']
                self.worker_loads[worker_id] = 0
                self.worker_speeds[worker_id] = worker.get('speed', 1.0)
                
            # Initialize task costs
            for task in self.tasks:
                self.task_costs[task['id']] = self._estimate_task_cost(task)
                
            # Create initial task segments
            self._create_task_segments()
            return True
        except Exception as e:
            print(f"Scheduler initialization failed: {str(e)}")
            return False
    
    def _estimate_task_cost(self, task: Dict[str, Any]) -> float:
        """Estimate time cost for a task based on computation complexity.
        
        Args:
            task: Task information including layer parameters
            
        Returns:
            float: Estimated time cost
        """
        params = task['parameters']
        # Basic cost estimation based on input size and operation type
        base_cost = params.get('input_size', 1) * params.get('output_size', 1)
        op_multiplier = {
            'conv2d': 2.0,
            'pool2d': 1.0,
            'linear': 1.5
        }.get(params.get('op_type', 'other'), 1.0)
        
        return base_cost * op_multiplier
    
    def _create_task_segments(self) -> None:
        """Create task segments based on synchronization points and costs.
        
        Divides tasks into segments considering task costs and worker speeds.
        """
        total_cost = sum(self.task_costs.values())
        avg_segment_cost = total_cost / (self.num_sync_points + 1)
        
        current_segment = []
        current_cost = 0
        self.task_segments = []
        
        for task in self.tasks:
            task_cost = self.task_costs[task['id']]
            current_segment.append(task)
            current_cost += task_cost
            
            if current_cost >= avg_segment_cost and len(self.task_segments) < self.num_sync_points:
                self.task_segments.append(current_segment)
                self.sync_points.append(len(self.tasks))
                current_segment = []
                current_cost = 0
        
        if current_segment:
            self.task_segments.append(current_segment)
    
    def assign_tasks(self) -> List[Dict[str, Any]]:
        """Assign tasks to workers based on performance-aware load balancing.
        
        Returns:
            List[Dict[str, Any]]: List of task assignments
        """
        assignments = []
        
        if not self.task_segments:
            return assignments
        
        current_segment = self.task_segments[0]
        if not current_segment:
            return assignments
        
        # Calculate normalized worker speeds
        total_speed = sum(self.worker_speeds.values())
        normalized_speeds = {w: s/total_speed for w, s in self.worker_speeds.items()}
        
        # Sort workers by current load/speed ratio
        sorted_workers = sorted(
            self.worker_loads.items(),
            key=lambda x: x[1]/self.worker_speeds[x[0]]
        )
        
        # Assign tasks proportionally to worker speeds
        worker_idx = 0
        for task in current_segment:
            worker_id = sorted_workers[worker_idx][0]
            task_cost = self.task_costs[task['id']]
            
            assignments.append({
                'task_id': task['id'],
                'worker_id': worker_id,
                'parameters': task['parameters']
            })
            
            # Update worker load
            self.worker_loads[worker_id] += task_cost
            
            # Move to next worker based on speed distribution
            worker_idx = (worker_idx + 1) % len(sorted_workers)
        
        return assignments
    
    def update_status(self, worker_id: str, task_id: str, status: str,
                      execution_time: Optional[float] = None) -> None:
        """Update task status and adjust worker performance metrics.
        
        Args:
            worker_id: Worker identifier
            task_id: Task identifier
            status: New task status
            execution_time: Actual task execution time if available
        """
        if status == 'completed':
            task_cost = self.task_costs[task_id]
            self.worker_loads[worker_id] -= task_cost
            
            # Update worker speed based on execution time
            if execution_time is not None and execution_time > 0:
                expected_time = task_cost
                speed_factor = expected_time / execution_time
                # Exponential moving average for speed update
                alpha = 0.3  # Update factor
                self.worker_speeds[worker_id] = (
                    alpha * speed_factor +
                    (1 - alpha) * self.worker_speeds[worker_id]
                )
    
    def handle_failure(self, worker_id: str, task_id: str) -> None:
        """Handle task failure by reassigning task and adjusting worker metrics.
        
        Args:
            worker_id: Failed worker identifier
            task_id: Failed task identifier
        """
        # Reduce worker speed and mark as potentially unstable
        self.worker_speeds[worker_id] *= 0.8  # Penalty factor
        self.worker_loads[worker_id] = float('inf')  # Prevent immediate reassignment
        
        # Find and requeue failed task
        failed_task = next((t for t in self.tasks if t['id'] == task_id), None)
        if failed_task:
            segment = next((seg for seg in self.task_segments if failed_task in seg), None)
            if segment:
                segment.append(failed_task)
    
    def get_sync_points(self) -> List[int]:
        """Get current synchronization points.
        
        Returns:
            List[int]: List of synchronization point indices
        """
        return self.sync_points
    
    def adjust_sync_points(self, performance_metrics: List[Dict[str, Any]]) -> None:
        """Adjust synchronization points based on runtime performance.
        
        Args:
            performance_metrics: List of worker performance measurements
        """
        if not performance_metrics:
            return
        
        # Calculate load imbalance
        execution_times = [m['execution_time'] for m in performance_metrics]
        max_time = max(execution_times)
        min_time = min(execution_times)
        imbalance = (max_time - min_time) / max_time if max_time > 0 else 0
        
        # Adjust number of sync points if imbalance is too high
        if imbalance > 0.3:  # Threshold for adjustment
            self.num_sync_points = min(self.num_sync_points + 1, len(self.tasks) // 2)
            self._create_task_segments()
        elif imbalance < 0.1 and self.num_sync_points > 1:
            self.num_sync_points -= 1
            self._create_task_segments()
    
    def free_nodes(self, completed_tasks: List[str]) -> List[str]:
        """Identify nodes that can be freed after task completion.
        
        Args:
            completed_tasks: List of completed task IDs
            
        Returns:
            List[str]: List of node IDs that can be freed
        """
        freeable_nodes = []
        
        for task_id in completed_tasks:
            task = next((t for t in self.tasks if t['id'] == task_id), None)
            if task:
                dependencies = task.get('dependencies', [])
                if all(dep in completed_tasks for dep in dependencies):
                    freeable_nodes.append(task_id)
        
        return freeable_nodes